import time

from ultralytics import YOLO
import cv2
import mypose

# 加载YOLO模型
model = YOLO('yolo11n-pose.pt')

# 读取图片
frame = cv2.imread('test3.png')
frameHeight, frameWidth = frame.shape[:2]

# 使用YOLO检测姿态关键点
results = model(frame)[0]
# 创建多个人的PoseKeypoints实例
poses = []
for i in range(len(results.boxes)):
    pose = mypose.PoseKeypoints(results, frameWidth, frameHeight, person_id=i)
    poses.append(pose)
    # 绘制每个人的骨架
    pose.PoseShow(frame, show_angles=True)
    time.sleep(0.5)

# 比较每对人的姿态相似度
for i in range(len(poses)):
    for j in range(i+1, len(poses)):
        mypose.compare_poses(frame, poses[i], poses[j])

# 显示结果
cv2.imshow('Multi-person Pose Comparison', frame)
cv2.waitKey(0)
cv2.destroyAllWindows()

# 保存结果图片
cv2.imwrite('result.png', frame)
